package com.star.langchain.learn.config;

import java.util.List;

import com.star.langchain.learn.aiservice.*;
import com.star.langchain.learn.myDefault.MyAiServices;
import com.star.langchain.learn.tools.BizTools;
import com.star.langchain.learn.tools.DateTools;
import dev.langchain4j.community.store.embedding.redis.RedisEmbeddingStore;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class MemoryConfig {
    /**
     * 智谱AI API密钥
     */
    private static final String API_KEY = "1166ca03dd9f407ea0ef76158c9b2f71.5u8R0mSPDIqynQKO";

    @Autowired
    private ChatMemoryStore chatMemoryStore;
    @Autowired
    private RedisEmbeddingStore redisEmbeddingStore;
    @Autowired
    private EmbeddingModel embeddingModel;
    @Autowired
    private ChatModel chatModel;
    @Autowired
    private StreamingChatModel streamingChatModel;
    @Autowired
    private McpToolProvider mcpToolProvider;




    @Bean
    public TokenStreamingService tokenStreamingService(){

        DateTools dateTools = new DateTools();
        BizTools bizTools = new BizTools();
        List<Object> tools = List.of(dateTools, bizTools);

        return MyAiServices.builder(TokenStreamingService.class)
                .chatModel(chatModel)
                .streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider())
                .tools(tools)
                .build();
    }


    @Bean
    public SimpleAiService simpleAiService() {
        DateTools dateTools = new DateTools();
        BizTools bizTools = new BizTools();
        List<Object> tools = List.of(dateTools, bizTools);
        return AiServices.builder(SimpleAiService.class)
                .chatModel(chatModel)
                .tools(tools)
                .chatMemoryProvider(chatMemoryProvider())
                .contentRetriever(contentRetriever())
                .inputGuardrailClasses(FirstGuard.class)
                .outputGuardrailClasses(SecondGuard.class)

                .toolProvider(mcpToolProvider)
                .build();
    }


    /**
     * 业务系统service
     *
     * @return
     */
    @Bean
    public BizSystemService bizSystemService() {
        DateTools dateTools = new DateTools();
        BizTools bizTools = new BizTools();
        List<Object> tools = List.of(dateTools, bizTools);
        return AiServices.builder(BizSystemService.class)
                .chatModel(chatModel)
                .chatMemoryProvider(chatMemoryProvider())
                .tools(tools)
                .build();
    }

    /**
     * mcpservice
     *
     * @return
     */
    @Bean
    public MCPService mcpService() {
        return AiServices.builder(MCPService.class)
                .chatModel(chatModel)
                .chatMemoryProvider(chatMemoryProvider())
                .toolProvider(mcpToolProvider)
                .build();
    }

    @Bean
    public RouterAgent routerAgent() {
        return AiServices.builder(RouterAgent.class)
                .chatModel(chatModel)
                .chatMemoryProvider(chatMemoryProvider())
                .tools(mcpService(), bizSystemService())
                .build();
    }

    @Bean
    ChatMemoryProvider chatMemoryProvider() {
        return memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .chatMemoryStore(chatMemoryStore)
                .maxMessages(30).build();
    }

    /**
     * 向量引入器
     *
     * @return 向量引入器
     */
    @Bean
    public EmbeddingStoreIngestor embeddingStoreIngestor() {
        return EmbeddingStoreIngestor.builder()
                // splitting each Document into TextSegments of 1000 tokens each, with a 200-token overlap
                .documentSplitter(DocumentSplitters.recursive(1000, 200))

//                // adding a name of the Document to each TextSegment to improve the quality of search
//                .textSegmentTransformer(textSegment -> TextSegment.from(
//                        textSegment.metadata().getString("file_name") + "\n" + textSegment.text(), textSegment.metadata()))
                .embeddingModel(embeddingModel).embeddingStore(redisEmbeddingStore).build();
    }

    /**
     * 内容检索器
     *
     * @return 内容检索器
     */
    @Bean
    public ContentRetriever contentRetriever() {
        return EmbeddingStoreContentRetriever.builder()
                .embeddingStore(redisEmbeddingStore)
                .embeddingModel(embeddingModel)
                .maxResults(3) // on each interaction we will retrieve the 2 most relevant
                // segments
                .minScore(0.8) // we want to retrieve segments at least somewhat similar to user query
                .build();
    }


}
